package com.own.component.store.redis.util;

import com.own.component.store.core.util.RateLimiterUtil;
import org.redisson.api.RateIntervalUnit;
import org.redisson.api.RateType;
import org.redisson.api.RedissonClient;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * RedisRateLimiterUtil
 *
 * @author chenxueli
 * @date 2024-04-16 15:00:00
 */
public class RedisRateLimiterUtil implements RateLimiterUtil {

    private static final Map<TimeUnit, RateIntervalUnit> MAP = new HashMap<>() {{
        put(TimeUnit.SECONDS, RateIntervalUnit.SECONDS);
        put(TimeUnit.MINUTES, RateIntervalUnit.MINUTES);
        put(TimeUnit.HOURS, RateIntervalUnit.HOURS);
        put(TimeUnit.DAYS, RateIntervalUnit.DAYS);
    }};
    private final RedissonClient redissonClient;

    public RedisRateLimiterUtil(RedissonClient redissonClient) {
        this.redissonClient = redissonClient;
    }

    /**
     * 获取RateLimiter
     *
     * @param name         名称
     * @param rate         速率
     * @param rateInterval 速率间隔
     * @param unit         单位
     */
    @Override
    public void create(String name, long rate, long rateInterval, TimeUnit unit) {
        var rateLimiter = redissonClient.getRateLimiter(name);
        rateLimiter.trySetRate(RateType.OVERALL, rate, rateInterval, MAP.get(unit));
    }

    /**
     * 尝试获取
     *
     * @param name 名称
     * @return 是否获取成功
     */
    @Override
    public boolean tryAcquire(String name) {
        return redissonClient.getRateLimiter(name).tryAcquire();
    }

}
